#include <bag_dealer/DataExtractor.hpp>

namespace BagDealer
{

    std::string ExtractorBase::timestampConverter(const std_msgs::Header &header)
    {
        auto secs = static_cast<long long int>(header.stamp.sec);
        auto nsecs = static_cast<long long int>(header.stamp.nsec);
        return std::to_string( static_cast<long long int>((secs * 1e6) +
                                                           static_cast<long long int>(nsecs / 1e3)));
    }

    std::string ExtractorBase::computeMD5(const std::string &input)
    {
        std::string md5;
        MD5_CTX ctx;
        MD5_Init(&ctx);
        MD5_Update(&ctx, input.c_str(), input.size());
        unsigned char digest[MD5_DIGEST_LENGTH];
        MD5_Final(digest, &ctx);
        char hex[35];
        memset(hex, 0, sizeof(hex));
        for (int i = 0; i < MD5_DIGEST_LENGTH; ++i) {
            sprintf(hex + i * 2, "%02x", digest[i]);
        }
        md5 = std::string(hex);
        return md5;
    }

    ros::Time ExtractorBase::getTimeFromStamp(const std_msgs::Header &ts)
    {
        auto sec = ts.stamp.sec;
        auto nsec = ts.stamp.nsec;
        return ros::Time(sec, nsec);
    }

    void SweepsExtractor::extract()
    {
        ROS_INFO_STREAM("Start to extract data from bag" << bagName);
        extractVelodyne();
//        extractCam("CAM_FRONT", "/usb_cam0/image_raw");
//        extractCam("CAM_LEFT", "/usb_cam1/image_raw");
//        extractCam("CAM_RIGHT", "/usb_cam2/image_raw");
        extractCam("CAM_FRONT",
                   nh_.param<std::string>("front_cam_topic", "/usb_cam0/image_raw"));
        extractCam("CAM_LEFT",
                   nh_.param<std::string>("left_cam_topic", "/usb_cam1/image_raw"));
        extractCam("CAM_RIGHT",
                   nh_.param<std::string>("right_cam_topic", "/usb_cam2/image_raw"));
        ROS_INFO_STREAM("DONE!");
    }

    void SweepsExtractor::extractVelodyne()
    {
        auto pcPath = sweepsPath / "LIDAR_TOP";
        if (!boost::filesystem::exists(pcPath))
        {
            boost::filesystem::create_directories(pcPath);
        }
        rosbag::View view(bag, rosbag::TypeQuery("sensor_msgs/PointCloud2"));
        pcl::PCDWriter pcWriter;


        ROS_INFO_STREAM("Writing POINTCLOUD sweeps!");
        BOOST_FOREACH(rosbag::MessageInstance const m, view)
        {
            auto pcMsgPtr = m.instantiate<sensor_msgs::PointCloud2>();
            pcl::PointCloud<pcl::PointXYZI> cloud;
            pcl::fromROSMsg(*pcMsgPtr, cloud);
            std::string timeCode = timestampConverter(pcMsgPtr->header);
            auto instancePath = pcPath / (bagNameNoSuffix + "__LIDAR_TOP__" + timeCode + ".pcd");

            pcWriter.write(instancePath.string(), cloud);
        }
    }

    void SweepsExtractor::extractCam(const std::string &camName, const std::string &topicName)
    {
        auto camPath = sweepsPath / camName;
        if (!boost::filesystem::exists(camPath))
        {
            boost::filesystem::create_directories(camPath);
        }
        rosbag::View view(bag, rosbag::TopicQuery(topicName));

        ROS_INFO_STREAM("Writing " << camName << " sweeps!");
        BOOST_FOREACH(rosbag::MessageInstance const m, view)
        {
            auto imgMsgPtr = m.instantiate<sensor_msgs::Image>();
            auto cv_ptr = cv_bridge::toCvCopy(imgMsgPtr, imgMsgPtr->encoding);
            cv::Mat img_cv = cv_ptr->image;
            cv::cvtColor(img_cv, img_cv, CV_RGB2BGR);
            std::string timeCode = timestampConverter(imgMsgPtr->header);
            auto instancePath = camPath /
                    (bagNameNoSuffix + "__" + camName + "__" + timeCode + ".jpg");
            std::vector<int> compression_params;
            compression_params.push_back(cv::IMWRITE_JPEG_QUALITY);
            compression_params.push_back(100);
            cv::imwrite(instancePath.string(), img_cv, compression_params);
        }
    }

    void SamplesExtractor::initSynchronizers()
    {
        message_filters::Subscriber<sensor_msgs::PointCloud2>
                lidar_sub(nh_,
                          nh_.param<std::string>("lidar_topic", "/velodyne_points"),
                          100);
        message_filters::Subscriber<sensor_msgs::Image>
                front_cam_sub(nh_,
                              nh_.param<std::string>("front_cam_topic", "/usb_cam0/image_raw"),
                              100);
        message_filters::Subscriber<sensor_msgs::Image>
                left_cam_sub(nh_,
                              nh_.param<std::string>("left_cam_topic", "/usb_cam1/image_raw"),
                              100);
        message_filters::Subscriber<sensor_msgs::Image>
                right_cam_sub(nh_,
                              nh_.param<std::string>("right_cam_topic", "/usb_cam2/image_raw"),
                              100);
        typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::Image,
                sensor_msgs::Image,
                sensor_msgs::Image,
                sensor_msgs::PointCloud2> myPolicy;
        message_filters::Synchronizer<myPolicy> synchronizer(myPolicy(100),
                                                             front_cam_sub,
                                                             left_cam_sub,
                                                             right_cam_sub,
                                                             lidar_sub);
        synchronizer.registerCallback(boost::bind(&SamplesExtractor::callback,
                                                  this,_1, _2,_3,_4));

        ros::spin();
    }

    bool SamplesExtractor::extractVelodyne(const sensor_msgs::PointCloud2ConstPtr &pcMsgPtr)
    {
        auto currentTime = getTimeFromStamp(pcMsgPtr->header);
        if (is_first || currentTime - keyFrameTime >=
        ros::Duration(1. / nh_.param<float>("hz", 2.5)))
        {
            is_first = false;
	    ROS_WARN_STREAM("time diff: " << currentTime - keyFrameTime);
            keyFrameTime = currentTime;
            auto pcPath = samplesPath / "LIDAR_TOP";
            if (!boost::filesystem::exists(pcPath)) {
                boost::filesystem::create_directories(pcPath);
            }
            pcl::PCDWriter pcWriter;

            pcl::PointCloud<pcl::PointXYZI> cloud;
            pcl::fromROSMsg(*pcMsgPtr, cloud);
            std::string timeCode = timestampConverter(pcMsgPtr->header);
            auto filename = bagNameNoSuffix + "__LIDAR_TOP__" + timeCode + ".pcd";
            auto instancePath = pcPath / filename;

            pcWriter.write(instancePath.string(), cloud);
            return true;
        }
        return false;
         
//        auto token = computeMD5(filename);
//        DS::SampleData* sdPtr = new DS::SampleData(token,
//                                                   "",
//                                                   "01fc3182f2699576b08fc6ad1c8661be",
//                                                   "samples" + filename,
//                                                   "pcd",
//                                                   0,0,
//                                                   timeCode,
//                                                   )
    }

    void SamplesExtractor::extractCam(const std::string &camName, const sensor_msgs::ImageConstPtr &imgMsgPtr)
    {
        auto camPath = samplesPath / camName;
        if (!boost::filesystem::exists(camPath))
        {
            boost::filesystem::create_directories(camPath);
        }
        auto cv_ptr = cv_bridge::toCvCopy(imgMsgPtr, imgMsgPtr->encoding);
        cv::Mat img_cv = cv_ptr->image;
        cv::cvtColor(img_cv, img_cv, CV_RGB2BGR);
        std::string timeCode = timestampConverter(imgMsgPtr->header);
        auto instancePath = camPath /
                            (bagNameNoSuffix + "__" + camName + "__" + timeCode + ".jpg");
        std::vector<int> compression_params;
        compression_params.push_back(cv::IMWRITE_JPEG_QUALITY);
        compression_params.push_back(100);
        cv::imwrite(instancePath.string(), img_cv, compression_params);
    }

    void SamplesExtractor::callback(const sensor_msgs::ImageConstPtr &frontImgMsg,
                                    const sensor_msgs::ImageConstPtr &leftImgMsg,
                                    const sensor_msgs::ImageConstPtr &rightImgMsg,
                                    const sensor_msgs::PointCloud2ConstPtr &pcMsg)
                                    {
                                        bool loading = extractVelodyne(pcMsg);
                                        if (loading)
                                        {
                                            extractCam("CAM_FRONT", frontImgMsg);
                                            extractCam("CAM_LEFT", leftImgMsg);
                                            extractCam("CAM_RIGHT", rightImgMsg);
                                        }

//                                        auto mainStamp = timestampConverter(pcMsg->header);
//                                        // sample.json here
//                                        auto codingStr = bagName +
//                                                timestampConverter(frontImgMsg->header) +
//                                                timestampConverter(leftImgMsg->header) +
//                                                timestampConverter(rightImgMsg->header) +
//                                                mainStamp;
//                                        auto token = computeMD5(codingStr);
//                                        auto prev = is_first ? "" : serializer.samples.back()->get_prev();
//                                        DS::Sample* sample = new DS::Sample(token,
//                                                          strtoll(mainStamp.c_str(), NULL, 10),
//                                                          computeMD5(bagName),
//                                                          prev);
//                                        is_first = false;
//                                        serializer.samples.back()->set_next(token);
//                                        serializer.samples.push_back(sample);
                                    }

    void SamplesExtractor::extract()
    {
        initSynchronizers();
    }

    SamplesExtractor::~SamplesExtractor()
    {
        if (bag.isOpen())
        {
            bag.close();
        }
    }
}
